import pickle
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from tqdm import tqdm

# 加载数据集
digits = load_digits()
X = digits.data
y = digits.target

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 初始化最佳参数
best_accuracy = 0.0
best_k = 1
best_knn_model = None

# 遍历K值(1-40)
k_list = range(1, 41)
accuracy_list = []

for k in tqdm(k_list, desc="训练KNN模型", unit="K值"):
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(X_train, y_train)
    y_pred = knn.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    accuracy_list.append(accuracy)
    
    if accuracy > best_accuracy:
        best_accuracy = accuracy
        best_k = k
        best_knn_model = knn

# 绘制准确率曲线
plt.figure(figsize=(10, 6))
plt.plot(k_list, accuracy_list, marker='o', color='b')
plt.axvline(x=best_k, color='r', linestyle='--', label=f'最优K值: {best_k}')
plt.text(best_k+0.5, best_accuracy, f'准确率: {best_accuracy:.3f}', color='r')
plt.xlabel('K值')
plt.ylabel('准确率')
plt.title('不同K值的KNN模型准确率')
plt.legend()
plt.grid(True, alpha=0.3)
plt.savefig('accuracy_plot.pdf')
plt.close()

# 保存最优模型
with open('best_knn_model.pkl', 'wb') as f:
    pickle.dump(best_knn_model, f)

print(f"训练完成! 最优K值: {best_k}, 准确率: {best_accuracy:.3f}")
print("已保存: accuracy_plot.pdf 和 best_knn_model.pkl")
    